from ultralytics import YOLO
from ultralytics.data.augment import LetterBox
import cv2

# 1. 读取原始图像
img = cv2.imread("D:/testpic/detect_20250927_103211_1758940331415.jpg")

# 2. 创建 LetterBox 实例（关键：去掉 scaleFill）
# auto=False: 不按 stride 对齐（与你 C++ 一致）
# center=True: 居中 padding（默认就是 True）
letterbox = LetterBox(new_shape=(320, 320), auto=False)

# 3. 应用 letterbox
letterboxed_img = letterbox(image=img)  # 输入 BGR HWC numpy array

# 4. 保存预处理图（用于和 C++ 的 preprocessed_letterboxed.jpg 对比）
cv2.imwrite("letterboxed_by_ultralytics.jpg", letterboxed_img)

# 5. 加载 ONNX 模型并推理
model = YOLO("D:/Code/helmet_detect_cpu/x64/Debug/best31n.onnx")
results = model(letterboxed_img, imgsz=320, verbose=False)

# 6. 在 320x320 图上画框（不反变换）
annotated = results[0].plot()
cv2.imwrite("ultralytics_on_320.jpg", annotated)

print("✅ Saved: letterboxed_by_ultralytics.jpg and ultralytics_on_320.jpg")